import json
import os

import util.common
from util.llm import http_request

MODEL_NAME = os.getenv('MODEL_NAME', "deepseek-chat")

SYSTEM_PROMPT = """你是一名应急预案专家，善于从事件类型和事件子类中提取关键词用于搜索相关应急预案文档。
我将给出一个事件描述，你需要理解这个事件，并从描述中提取关键词用于搜索。
提取的关键词可以是多个，关键词使用json array表示。
输出示例：
如果需要输出的内容为"网络安全事件"、"数据勒索"、"应急响应"，则返回的内容为：["网络安全事件","数据勒索","应急响应"]"""

PROMPT = """以下我将给出一个事件描述，请帮我提取关键词用于搜索应急预案文档。结果用json array表示
事件描述：{}"""


async def query_rewrite(query):
    """
    使用大模型进行查询重写（原有方法）
    """
    print(f"🔍 query_rewrite开始处理，输入查询: {query}")
    
    try:
        result = await http_request(
            system_prompt=SYSTEM_PROMPT,
            history=[],
            prompt=PROMPT.format(query),
            model=MODEL_NAME
        )
        print(f"✅ LLM查询重写成功，原始结果: {result}")
        
        result = util.common.remove_markdown_code_block(result)
        print(f"🔍 移除markdown代码块后: {result}")
        
        result = json.loads(result)
        print(f"🔍 JSON解析成功: {result}")
        
        result.append(query)
        result = list(set(result))
        print(f"✅ 查询重写完成，最终结果: {result}")
        
    except Exception as e:
        print(f"❌ query_rewrite失败: {e}")
        import traceback
        traceback.print_exc()
        result = [query, ]
        print(f"🔍 使用原始查询作为备选: {result}")
    
    return result


async def query_rewrite_retrieval_only(query):
    """
    纯检索模式的查询重写，不使用大模型（新增方法）
    """
    print(f"🔍 query_rewrite_retrieval_only开始处理，输入查询: {query}")
    
    try:
        # 直接基于输入信息提取关键词
        print(f"🔍 使用纯检索模式，直接关键词提取")
        
        # 基于查询内容直接提取关键词
        keywords = []
        
        # 添加原始查询
        keywords.append(query)
        
        # 从查询中提取可能的关键词
        if "网络安全" in query:
            keywords.extend(["网络安全", "网络攻击", "数据泄露"])
        if "数据" in query:
            keywords.extend(["数据安全", "数据保护", "数据泄露"])
        if "应急" in query:
            keywords.extend(["应急响应", "应急预案", "应急处置"])
        if "事件" in query:
            keywords.extend(["安全事件", "突发事件", "应急事件"])
        if "攻击" in query:
            keywords.extend(["网络攻击", "恶意攻击", "安全攻击"])
        if "泄露" in query:
            keywords.extend(["数据泄露", "信息泄露", "隐私泄露"])
        if "勒索" in query:
            keywords.extend(["勒索软件", "数据勒索", "勒索攻击"])
        
        # 去重
        keywords = list(set(keywords))
        
        print(f"✅ 纯检索模式关键词提取完成: {keywords}")
        
    except Exception as e:
        print(f"❌ query_rewrite_retrieval_only失败: {e}")
        import traceback
        traceback.print_exc()
        keywords = [query, ]
        print(f"🔍 使用原始查询作为备选: {keywords}")
    
    return keywords 